// SPDX-License-Identifier: GPL-2.0

#include <linux/module.h>
#include <linux/file.h>
#include <linux/fs.h>
#include <linux/idr.h>
#include <linux/mutex.h>
#include <linux/slab.h>
#include <linux/miscdevice.h>
#include <linux/uaccess.h>

#include <rdma/ib.h>
#include <rdma/ib_user_cm.h>
#include <trace/hooks/rdma_pangea.h>

static DEFINE_IDR(fd_idr);
static DEFINE_MUTEX(mut);
static LIST_HEAD(fd_list);

#ifndef CONFIG_MODULE_UPGRADE
#define INVALID_HANDLE (-1)
#define INVALID_FD (-1)

struct fd_container {
	struct list_head entry;
	__u32 port;
	__u32 handle;
	struct file *file;
};

static int __ib_ext_fd_hold(int fd, __u32 port)
{
	int handle;
	struct file *file;
	struct fd_container *container;

	if (fd <= 0) {
		ofed_log_print(PANGEA_ERR, "Bad fd(%d).", fd);
		goto err;
	}

	container = kzalloc(sizeof(*container), GFP_KERNEL);
	if (!container) {
		ofed_log_print(PANGEA_ERR, "Get free memory failed.");
		goto err;
	}

	file = fget(fd);
	if (!file) {
		ofed_log_print(PANGEA_ERR, "Get file failed, fd(%d).", fd);
		goto err_free;
	}

	container->file = file;
	container->port = port;

	handle = idr_alloc_cyclic(&fd_idr, container, 1, 0, GFP_KERNEL);
	if (handle <= 0) {
		ofed_log_print(PANGEA_ERR, "Invalid idr handle(%d).", handle);
		goto err_put;
	}

	container->handle = handle;
	list_add_tail(&container->entry, &fd_list);
	return handle;
err_put:
	fput(file);
err_free:
	kfree(container);
err:
	return INVALID_HANDLE;
}
#endif

static ssize_t ib_ext_fd_hold(struct file *filp,
	const char __user *inbuf, int in_len)
{
	int i, j;
	struct fd_container *container;
	struct ib_ext_hold_fd_resp resp = { 0 };
	struct ib_ext_hold_fd cmd = { 0 };

	if (copy_from_user(&cmd, inbuf, sizeof(cmd))) {
		ofed_log_print(PANGEA_ERR, "Copy from user failed.");
		return -EFAULT;
	}

	if (cmd.fd_num > MAX_FD_NUMBER) {
		ofed_log_print(PANGEA_ERR, "Invalid fd number(%d).", cmd.fd_num);
		return -EINVAL;
	}

	mutex_lock(&mut);
	for (i = 0; i < cmd.fd_num; i++) {
#ifndef CONFIG_MODULE_UPGRADE
		resp.handle[i] = __ib_ext_fd_hold(cmd.fd[i], cmd.port);
#endif
		if (resp.handle[i] == INVALID_HANDLE)
			goto err;
	}
	mutex_unlock(&mut);

	if (copy_to_user(u64_to_user_ptr(cmd.response), &resp, sizeof(resp))) {
		ofed_log_print(PANGEA_ERR, "Copy to user failed, user response(0x%llx).",
			cmd.response);
		mutex_lock(&mut);
		goto err;
	}

	return 0;
err:
	for (j = 0; j < i; j++) {
		container = (struct fd_container *)idr_find(&fd_idr, resp.handle[j]);
		if (!container) {
			ofed_log_print(PANGEA_ERR, "Can not find file by handle(%d).",
				resp.handle[j]);
			continue;
		}
		fput(container->file);
		idr_remove(&fd_idr, resp.handle[j]);
		list_del(&container->entry);
		kfree(container);
	}
	mutex_unlock(&mut);

	return -EINVAL;
}

static void __ib_ext_recover_fd_error(struct fd_container **containers,
	unsigned int fd_num)
{
	unsigned int i;
	struct file *file;
	struct fd_container *container;

	for (i = 0; i < fd_num; i++) {
		container = containers[i];
		file = container->file;
		fput(file);
		idr_remove(&fd_idr, container->handle);
		list_del(&container->entry);
		kfree(container);
	}
}

static int ib_ext_get_unused_fds(int *fd, unsigned int fd_num)
{
	int ret;

	ret = get_unused_fd_flags(O_CLOEXEC);
	if (ret < 0) {
		ofed_log_print(PANGEA_ERR, "Get unused fd failed, ret(%d).", ret);
		return ret;
	}

	fd[0] = ret;

	if (fd_num > 1) {
		ret = get_unused_fd_flags(O_CLOEXEC);
		if (ret < 0) {
			ofed_log_print(PANGEA_ERR, "Get unused fd failed, ret(%d).", ret);
			goto err;
		}
	}
	fd[1] = ret;

	return 0;
err:
	put_unused_fd(fd[0]);
	return ret;
}

#ifndef CONFIG_MODULE_UPGRADE
static void ib_ext_fd_install_file(int *fd, struct fd_container **containers,
	unsigned int fd_num)
{
	unsigned int i;
	struct file *file = NULL;

	for (i = 0; i < fd_num; i++) {
		file = containers[i]->file;
		get_file(file);
		fd_install(fd[i], file);
		if (file->f_cred)
			put_cred(file->f_cred);
		file->f_cred = get_cred(current_cred());
	}
}
#endif

static ssize_t ib_ext_fd_recover(struct file *filp,
	const char __user *inbuf, int in_len)
{
	int i, ret;
	struct ib_ext_recover_fd cmd = { 0 };
	struct ib_ext_recover_fd_resp resp = { 0 };
	struct fd_container *container[2] = {NULL, NULL};
	int fd[2] = {-1, -1};

	if (copy_from_user(&cmd, inbuf, sizeof(cmd))) {
		ofed_log_print(PANGEA_ERR, "Copy from user failed.");
		return -EFAULT;
	}

	if (cmd.fd_num > MAX_FD_NUMBER) {
		ofed_log_print(PANGEA_ERR, "Invalid fd number(%d).", cmd.fd_num);
		return -EINVAL;
	}

	ret = ib_ext_get_unused_fds(fd, cmd.fd_num);
	if (ret)
		return ret;

	mutex_lock(&mut);
	for (i = 0; i < cmd.fd_num; i++) {
		if (cmd.handle[i] <= 0) {
			ofed_log_print(PANGEA_ERR, "Invalid handle(%d).", cmd.handle[i]);
			goto err;
		}

		container[i] = (struct fd_container *)idr_find(&fd_idr, cmd.handle[i]);
		if (container[i] == NULL) {
			ofed_log_print(PANGEA_ERR, "Can not find file by handle(%d).",
				cmd.handle[i]);
			goto err;
		}
	}

	resp.fd[0] = fd[0];
	resp.fd[1] = fd[1];
	if (copy_to_user(u64_to_user_ptr(cmd.response), &resp, sizeof(resp))) {
		ofed_log_print(PANGEA_ERR, "Copy to user failed, user response(0x%llx).",
			cmd.response);
		ret = -EFAULT;
	}

	ib_ext_fd_install_file(fd, container, cmd.fd_num);
	mutex_unlock(&mut);
	return 0;
err:
	__ib_ext_recover_fd_error(container, i);

	put_unused_fd(fd[0]);
	if (fd[1] > 0)
		put_unused_fd(fd[1]);
	mutex_unlock(&mut);
	return -EINVAL;
}

static ssize_t ib_ext_fd_remove(struct file *filp,
	const char __user *inbuf, int in_len)
{
	int i;
	struct file *file;
	struct fd_container *container;
	struct ib_ext_remove_fd cmd = { 0 };
	struct ib_ext_remove_fd_resp resp = { 0 };

	if (copy_from_user(&cmd, inbuf, sizeof(cmd))) {
		ofed_log_print(PANGEA_ERR, "Copy from user failed.");
		return -EFAULT;
	}

	if (cmd.fd_num > MAX_FD_NUMBER) {
		ofed_log_print(PANGEA_ERR, "Invalid fd number(%d).", cmd.fd_num);
		return -EINVAL;
	}

	mutex_lock(&mut);
	for (i = 0; i < cmd.fd_num; i++) {
		container = (struct fd_container *)idr_find(&fd_idr, cmd.handle[i]);
		resp.handle[i] = INVALID_HANDLE;
		if (!container) {
			ofed_log_print(PANGEA_ERR, "Can not find file by handle(%d).",
				cmd.handle[i]);
			continue;
		}
		file = container->file;
		fput(file);
		idr_remove(&fd_idr, cmd.handle[i]);
		list_del(&container->entry);
		kfree(container);
	}
	mutex_unlock(&mut);

	if (copy_to_user(u64_to_user_ptr(cmd.response), &resp, sizeof(resp))) {
		ofed_log_print(PANGEA_ERR, "Copy to user failed, user response(0x%llx).",
			cmd.response);
		return -EFAULT;
	}

	return 0;
}

static ssize_t ib_ext_fd_clean(struct file *filp,
	const char __user *inbuf, int in_len)
{
	struct fd_container *container;
	struct fd_container *tmp;
	struct file *file;
	struct ib_ext_clean_fd cmd = { 0 };

	if (copy_from_user(&cmd, inbuf, sizeof(cmd))) {
		ofed_log_print(PANGEA_ERR, "Copy from user failed.");
		return -EFAULT;
	}

	mutex_lock(&mut);
	list_for_each_entry_safe(container, tmp, &fd_list, entry) {
		if (container->port != cmd.port)
			continue;

		file = container->file;
#ifndef CONFIG_MODULE_UPGRADE
		fput(file);
#endif

		idr_remove(&fd_idr, container->handle);
		list_del(&container->entry);
		kfree(container);
	}
	mutex_unlock(&mut);

	return 0;
}

static ssize_t ib_ext_fd_rsvd1(struct file *filp,
	const char __user *inbuf, int in_len)
{
	return -EOPNOTSUPP;
}

static ssize_t ib_ext_fd_rsvd2(struct file *filp,
	const char __user *inbuf, int in_len)
{
	return -EOPNOTSUPP;
}

static ssize_t ib_ext_fd_rsvd3(struct file *filp,
	const char __user *inbuf, int in_len)
{
	return -EOPNOTSUPP;
}

static ssize_t ib_ext_fd_rsvd4(struct file *filp,
	const char __user *inbuf, int in_len)
{
	return -EOPNOTSUPP;
}

static ssize_t (*fd_cmd_table[])(struct file *file,
	const char __user *inbuf,
	int in_len) = {
	[IB_EXT_FD_CMD_HOLD] = ib_ext_fd_hold,
	[IB_EXT_FD_CMD_REMOVE] = ib_ext_fd_remove,
	[IB_EXT_FD_CMD_RECOVER] = ib_ext_fd_recover,
	[IB_EXT_FD_CMD_CLEAN] = ib_ext_fd_clean,
	[IB_EXT_FD_CMD_RSVD1] = ib_ext_fd_rsvd1,
	[IB_EXT_FD_CMD_RSVD2] = ib_ext_fd_rsvd2,
	[IB_EXT_FD_CMD_RSVD3] = ib_ext_fd_rsvd3,
	[IB_EXT_FD_CMD_RSVD4] = ib_ext_fd_rsvd4,
};

static ssize_t ib_ext_fd_write(struct file *filp, const char __user *buf,
	size_t len, loff_t *pos)
{
	int ret;
	struct ib_ext_fd_hdr hdr = { 0 };

	if (!ib_safe_file_access(filp)) {
		ofed_log_print(PANGEA_ERR, "Process %d (%s) changed security contexts.",
			task_tgid_vnr(current), current->comm);
		return -EACCES;
	}

	if (len < sizeof(hdr)) {
		ofed_log_print(PANGEA_ERR, "Invalid hdr len(%zu).", len);
		return -EINVAL;
	}

	if (copy_from_user(&hdr, buf, sizeof(hdr))) {
		ofed_log_print(PANGEA_ERR, "Copy from user failed.");
		return -EFAULT;
	}

	if (hdr.cmd >= ARRAY_SIZE(fd_cmd_table)) {
		ofed_log_print(PANGEA_ERR, "Invalid cmd(%d).", hdr.cmd);
		return -EINVAL;
	}

	if ((hdr.in + sizeof(hdr)) > len) {
		ofed_log_print(PANGEA_ERR, "Invalid len(%zu), in(%u).", len, hdr.in);
		return -EINVAL;
	}

	ret = fd_cmd_table[hdr.cmd](filp, buf + sizeof(hdr), hdr.in);
	if (ret == 0)
		ret = len;

	return ret;
}

static const struct file_operations ext_fd_fops = {
	.owner	= THIS_MODULE,
	.open	= nonseekable_open,
	.write	= ib_ext_fd_write,
	.llseek	= no_llseek,
};

static struct miscdevice ext_fd_misc = {
	.minor		= MISC_DYNAMIC_MINOR,
	.name		= "ucm_fd",
	.nodename	= "infiniband/ucm_fd",
	.mode       = 0666,
	.fops		= &ext_fd_fops,
};

int ib_ext_fd_init(void)
{
	return misc_register(&ext_fd_misc);
}

void ib_ext_fd_exit(void)
{
	misc_deregister(&ext_fd_misc);
}

